#!/usr/bin/env python
# vim: set fileencoding=utf-8 :
#
# This file is part of Coruja-scripts
#
# Coruja-scripts is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, version 3 of the License.
#
# Coruja-scripts is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Coruja-scripts.  If not, see <http://www.gnu.org/licenses/>.
#
# Copyright 2011 Grupo Falabrasil - http://www.laps.ufpa.br/falabrasil
#
# Author 2011: Pedro Batista - pedosb@gmail.com

import argparse
import re
import os
from energy_estimation import mkdir_p
import scipy.io.wavfile as wf
import numpy as np
from config import *

def add_noise(signal, snr, out):
	fs, data = wf.read(signal)
	noise = np.random.randn(len(data))
	signal_power = sum(data**2)
	d_noise_power = signal_power / (10 ** (snr/10.))
	noise *= np.sqrt(d_noise_power / len(noise))
	print 10 * np.log10(sum(data**2)/sum(noise**2))
	data += noise
	wf.write(out, fs, data)

def add_noise_list(wav_list, snr, mask, replacement):
	for line in open(wav_list):
		line = line.strip()
		out_file = re.sub(mask, replacement, line)
		d = os.path.dirname(out_file)
		if d:
			mkdir_p(d)
		add_noise(line, snr,
				out_file)

if __name__=='__main__':
	parser = argparse.ArgumentParser(description='Add aditive gausian noise to a wave signal')
	parser.add_argument('-s', '--signal',
			help='The wav file that contains the signal.')
	parser.add_argument('-l', '--input-file-list',
			help='List of WAV input files name')
	parser.add_argument('-o', '--output-signal',
			help='The wav file that contains the signal plus the noise.')
	parser.add_argument('-n', '--snr', required=True, type=float,
			default=config.getfloat('noise', 'snr'),
			help='The resulting signal to noise ratio.')
	parser.add_argument('-m', '--input-file-mask',
			default=config.get('noise', 'noise_input_mask'),
			help='Mask to match the input file')
	parser.add_argument('-r', '--replacement',
			default=config.get('noise', 'noise_input_replacement'),
			help='Replacement from input-file-mask to produce output file')
	args = parser.parse_args()
	if args.snr != config.getfloat('noise', 'snr'):
		config.set('noise', 'snr', str(args.snr))
		args.input_file_mask = config.get('noise', 'noise_input_mask')
		args.replacement = config.get('noise', 'noise_input_replacement')
	if args.input_file_list:
		add_noise_list(args.input_file_list,
				args.snr,
				args.input_file_mask,
				args.replacement)
	else:
		add_noise(args.signal, args.snr, args.output_signal)
